//--------------------------------------------------------------------------------------
// File: ComputeShaderSort11.hlsl
//
// This file contains the compute shaders to perform GPU sorting using DirectX 11.
// 
// Copyright (c) Microsoft Corporation. All rights reserved.
//--------------------------------------------------------------------------------------

#ifndef TYPE
#define TYPE float
#endif
#ifndef BITONIC_BLOCK_SIZE
#define BITONIC_BLOCK_SIZE 1024
#endif
#ifndef TRANSPOSE_BLOCK_SIZE
#define TRANSPOSE_BLOCK_SIZE 32
#endif

//--------------------------------------------------------------------------------------
// Constant Buffers
//--------------------------------------------------------------------------------------
cbuffer CB : register (b0)
{
    uint g_iLevel;
    uint g_iLevelMask;
    uint g_iWidth;
    uint g_iHeight;
};

StructuredBuffer<TYPE> Input : register (t0);
RWStructuredBuffer<TYPE> Data : register (u0);

//--------------------------------------------------------------------------------------
// Bitonic Sort Compute Shader
//--------------------------------------------------------------------------------------


groupshared TYPE shared_data[BITONIC_BLOCK_SIZE];

[numthreads(BITONIC_BLOCK_SIZE / 2, 1, 1)]
void BitonicSort( uint3 Gid : SV_GroupID, 
                  uint3 DTid : SV_DispatchThreadID, 
                  uint3 GTid : SV_GroupThreadID, 
                  uint GI : SV_GroupIndex )
{
    // Load shared data
	shared_data[GI*2] = Data[DTid.x*2];
	shared_data[GI*2+1] = Data[DTid.x*2+1];
//	GroupMemoryBarrierWithGroupSync();
    
    // Sort the shared data
	bool asc = !(bool)(g_iLevelMask & (GI * 2 + Gid.x * BITONIC_BLOCK_SIZE));
	uint j = g_iLevel >> 1;
    for (; j > 64; j >>= 1) {
//		TYPE result = ((shared_data[GI & ~j].x <= shared_data[GI | j].x) == (bool)(g_iLevelMask & DTid.x)) ? shared_data[GI ^ j] : shared_data[GI];
//        GroupMemoryBarrierWithGroupSync();
//        shared_data[GI] = result;
        GroupMemoryBarrierWithGroupSync();
		float idx1 = GI / j * j * 2 + GI % j;
		float idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
    }
	GroupMemoryBarrierWithGroupSync();
	if (j == 64) {
 		float idx1 = GI / j * j * 2 + GI % j;
		float idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
	} else if (j == 32) {
 		float idx1 = GI / j * j * 2 + GI % j;
		float idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
	} else if (j == 16) {
 		float idx1 = GI / j * j * 2 + GI % j;
		float idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
	} else if (j == 8) {
 		float idx1 = GI / j * j * 2 + GI % j;
		float idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
	} else if (j == 4) {
 		float idx1 = GI / j * j * 2 + GI % j;
		float idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
	} else if (j == 2) {
 		float idx1 = GI / j * j * 2 + GI % j;
		float idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
		j >>= 1;
 		idx1 = GI / j * j * 2 + GI % j;
		idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
	} else {
 		float idx1 = GI / j * j * 2 + GI % j;
		float idx2 = GI / j * j * 2 + GI % j + j;
		if ((shared_data[idx1].x <= shared_data[idx2].x) != asc) {
			TYPE temp = shared_data[idx1];
			shared_data[idx1] = shared_data[idx2];
			shared_data[idx2] = temp;
		}
        GroupMemoryBarrierWithGroupSync();
	}
   
    // Store shared data
	//Data[DTid.x] = shared_data[GI];
    Data[DTid.x*2] = shared_data[GI*2];
	Data[DTid.x*2+1] = shared_data[GI*2+1];
}

//--------------------------------------------------------------------------------------
// Matrix Transpose Compute Shader
//--------------------------------------------------------------------------------------
groupshared TYPE transpose_shared_data[TRANSPOSE_BLOCK_SIZE * TRANSPOSE_BLOCK_SIZE];

[numthreads(TRANSPOSE_BLOCK_SIZE, TRANSPOSE_BLOCK_SIZE, 1)]
void MatrixTranspose (uint3 Gid : SV_GroupID, 
                      uint3 DTid : SV_DispatchThreadID, 
                      uint3 GTid : SV_GroupThreadID, 
                      uint GI : SV_GroupIndex)
{
    transpose_shared_data[GI] = Input[DTid.y * g_iWidth + DTid.x];
    GroupMemoryBarrierWithGroupSync ();
    uint2 XY = DTid.yx - GTid.yx + GTid.xy;
    Data[XY.y * g_iHeight + XY.x] = transpose_shared_data[GTid.x * TRANSPOSE_BLOCK_SIZE + GTid.y];
}
